import bpy, os
import numpy as np
from tool.cam import Cam
from math import radians
import random
from remote_run_everything import Common
from tool.ext_convert import from_object_to_uv

class Cylinder:
    def __init__(self):
        self.root = "D://ai/my_blender"
        self.d = "D://ai/my_blender/data"
        self.d2 = "D://ai/my_blender/data2"
        os.makedirs(self.d, exist_ok=True)
        os.makedirs(self.d2, exist_ok=True)
        self.scale = 0.5
        self.r = 0.6 * self.scale
        self.h = 1.2 * self.scale
        self.nu = 13
        self.cam = Cam()
        self.K, self.W, self.H = self.cam.intrinsic()
        self.ex = self.cam.extrinsic()
        self.data = {}

    def count_data(self):
        l = []
        for _, v in self.data.items():
            l = v + l
        total = sum(l)
        return total

    def clear(self):
        objects = bpy.data.objects
        for k in objects.keys():
            if k != "Camera" and k != "Light":
                objects.remove(objects[k])
        bpy.ops.outliner.orphans_purge(do_recursive=True)

    def add_cyl(self, loc):
        bpy.ops.mesh.primitive_cylinder_add(
            radius=self.r,
            depth=self.h,
            rotation=(np.pi / 2, 0, 0),
            location=loc,
        )
        return bpy.context.active_object

    def get_data(self, level, full):
        if full:
            l = [1 for i in range(self.nu)]
        else:
            l = [random.choice([0, 1]) for i in range(self.nu)]
        if level >= 2: l[-1] = 0
        if level >= 3: l[-2] = 0
        if level >= 2:
            l0 = self.data.get(level - 1, [])
            if len(l0) == 0:
                self.data[level] = []
                return
            for pos, v in enumerate(l):
                if pos >= self.nu - 1: break
                left = l0[pos]
                right = l0[pos + 1]
                if left == 0 or right == 0: l[pos] = 0
        self.data[level] = l
        return l

    def add_row(self, y=0):
        locs=[]
        step = 2 * self.r
        for level, data in self.data.items():
            start = -self.r * self.nu + (level - 1) * self.r
            z = (level - 1) * 1.7 * self.r
            for ix, v in enumerate(data):
                if v == 0:
                    # print("not", level, ix, v)
                    continue
                loc = [start + ix * step, y, z]
                ob = self.add_cyl(loc)
                # from_object_to_uv(ob)
                locs.append(loc)
                # self.cam.wordloc2uv(loc)
        return locs


    def run(self):
        self.clear()
        self.get_data(1, True)
        # self.get_data(2,True)
        # self.get_data(3, False)
        self.add_row()
        # total = self.count_data()
        # nu = Common().prefix_zero(6, 1)
        path = f"{self.d}/test.png"
        self.cam.render(path)


if __name__ == '__main__':
    cy = Cylinder()
    cy.run()
